import pickle
import os
import argparse
from pprint import pprint

from utils.metric import METRIC_REGISTRY
from utils.action import ACTION_REGISTRY

def parse_args():
    parser = argparse.ArgumentParser(description='evaluate net')
    parser.add_argument("--model", type=str, help='model name', required=True)
    parser.add_argument("--datasets", type=str, help='train datasets name', required=True)
    parser.add_argument("--val-datasets", type=str, help='val datasets name', required=True)
    parser.add_argument("--test-datasets", type=str, help='datasets for sampling', required=True)
    parser.add_argument("--checkpoint", type=str, help="checkpoint iter, can be 'last', 'best', iternumber", default='best')
    parser.add_argument("--nms-threshold", type=float, help='nms threshold', default=0.5)
    parser.add_argument("--iou-cover-threshold", type=float, help='iou cover threshold', default=0.5)
    parser.add_argument("--proposals-cover-threshold", type=float, help='proposals cover threshold', default=0.7)
    parser.add_argument("--gt-boxes-cover-threshold", type=float, help='gt boxes cover threshold', default=0.7)
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()

    train_datasets = ["{}-train".format(dataset) for dataset in args.datasets.split(',')]
    val_datasets = ["{}-test".format(dataset) for dataset in args.val_datasets.split(',')]
    test_datasets = ["{}-test".format(dataset) for dataset in args.test_datasets.split(',')]

    evaluate_path = os.path.join('output','evaluate', '{}/[train]-{}/[val]-{}/[test]-{}-[checkpoint]-{}/[nms]-{}-[evaluate]-{}-{}-{}.pkl'.format(args.model, '&'.join(train_datasets), '&'.join(val_datasets), '&'.join(test_datasets), args.checkpoint, args.nms_threshold, args.iou_cover_threshold, args.proposals_cover_threshold, args.gt_boxes_cover_threshold))
    with open(evaluate_path, 'rb') as f:
        evaluate_dict_list = pickle.load(f)
    evaluate_dict_list.sort(key=lambda evaluate_dict:evaluate_dict['threshold'])

    output_dir = os.path.join('output','analysis', '{}/[train]-{}/[val]-{}/[test]-{}-[checkpoint]-{}/[nms]-{}-[evaluate]-{}-{}-{}'.format(args.model, '&'.join(train_datasets), '&'.join(val_datasets), '&'.join(test_datasets), args.checkpoint, args.nms_threshold, args.iou_cover_threshold, args.proposals_cover_threshold, args.gt_boxes_cover_threshold))
    os.makedirs(output_dir, exist_ok=True)

    # do all action
    for action_name in ACTION_REGISTRY.registered_names():
        ACTION_REGISTRY.get(action_name)(evaluate_dict_list, output_dir)

    # do all metric
    output_path = os.path.join(output_dir, 'best_threshold.txt')
    with open(output_path, 'w', encoding='utf-8') as f:
        for metric_name in METRIC_REGISTRY.registered_names():
            print("================================", file=f)
            print("---METRIC---: {}".format(metric_name), file=f)
            print("================================", file=f)

            result_dict = METRIC_REGISTRY.get(metric_name)(evaluate_dict_list)
            for key, value in result_dict.items():
                print("------------ {} ------------".format(key), file=f)
                pprint(value, stream=f)

    print("analysis files save at {}".format(output_dir))
    print("done")